"""
异步数据库连接池管理器
提供高性能、可监控、自适应的PostgreSQL连接池实现
"""
import asyncio
import time
from typing import Dict, Any, Optional, List, Callable, AsyncGenerator
import logging
from datetime import datetime, timedelta

from sqlalchemy.ext.asyncio import (
    AsyncEngine,
    AsyncSession,
    async_sessionmaker,
    create_async_engine,
    AsyncConnection,
)
from sqlalchemy.sql import text

from app.core.config import settings
from utils.logger import logger


class AsyncDBPoolManager:
    """增强版异步数据库连接池管理器"""
    
    def __init__(self):
        self.engine: Optional[AsyncEngine] = None
        self.async_session_factory: Optional[async_sessionmaker[AsyncSession]] = None
        self._pool_stats: Dict[str, Any] = {
            "active_connections": 0,
            "idle_connections": 0,
            "total_connections": 0,
            "connection_errors": 0,
            "connection_successes": 0,
            "last_stats_update": datetime.now(),
        }
        self._health_check_interval = 60  # 秒
        self._health_check_task: Optional[asyncio.Task] = None
        self._pool_adjustment_enabled = True
        self._last_pool_size = settings.DATABASE.pool_size
        self._connection_metrics: Dict[str, List[Dict[str, Any]]] = {}
        self._retry_policy = {
            "max_retries": 3,
            "retry_delay": 0.5,
            "backoff_factor": 2.0,
        }
    
    async def initialize(self):
        """
        初始化连接池
        """
        try:
            # 构建数据库URL
            db_url = f"postgresql+asyncpg://{settings.DATABASE.username}:{settings.DATABASE.password}@{settings.DATABASE.host}:{settings.DATABASE.port}/{settings.DATABASE.database}"
            
            logger.info(f"正在初始化数据库连接池: {settings.DATABASE.host}:{settings.DATABASE.port}")
            logger.info(f"连接池配置: pool_size={settings.DATABASE.pool_size}, max_overflow={settings.DATABASE.max_overflow}, pool_timeout={settings.DATABASE.pool_timeout}秒")
            
            # 创建优化的异步引擎
            self.engine = create_async_engine(
                db_url,
                echo=settings.APP_DEBUG,  # 调试模式下打印SQL
                pool_size=settings.DATABASE.pool_size,
                max_overflow=settings.DATABASE.max_overflow,
                pool_timeout=settings.DATABASE.pool_timeout,
                pool_recycle=settings.DATABASE.pool_recycle,
                pool_pre_ping=True,  # 连接借出前进行ping测试
                pool_reset_on_return="commit",  # 连接归还时重置
                future=True,
                connect_args={
                    "command_timeout": 30,  # 命令超时时间
                    "timeout": 10,  # 连接超时时间
                    "ssl": None,  # 根据环境配置SSL
                },
            )
            
            # 创建异步会话工厂
            self.async_session_factory = async_sessionmaker(
                self.engine,
                class_=AsyncSession,
                expire_on_commit=False,
                autocommit=False,
                autoflush=False,
                future=True,
            )
            
            logger.info("数据库连接池初始化成功")
            
            # 启动健康检查任务
            self._start_health_check()
            
            return self.engine
            
        except Exception as e:
            logger.error(f"数据库连接池初始化失败: {str(e)}")
            raise
    
    async def close(self):
        """
        关闭连接池
        """
        # 停止健康检查任务
        if self._health_check_task:
            self._health_check_task.cancel()
            try:
                await self._health_check_task
            except asyncio.CancelledError:
                pass
        
        # 关闭引擎
        if self.engine:
            await self.engine.dispose()
            logger.info("数据库连接池已关闭")
    
    async def get_db_session(self) -> AsyncGenerator[AsyncSession, None]:
        """
        获取数据库会话（带重试机制）
        """
        if not self.async_session_factory:
            raise RuntimeError("数据库会话工厂未初始化")
        
        retries = 0
        while retries <= self._retry_policy["max_retries"]:
            try:
                start_time = time.time()
                async with self.async_session_factory() as session:
                    try:
                        self._increment_connection_count()
                        yield session
                        # 记录成功的连接使用情况
                        self._pool_stats["connection_successes"] += 1
                        self._record_connection_metric("success", time.time() - start_time)
                    finally:
                        self._decrement_connection_count()
                        await session.close()
                break
                
            except Exception as e:
                retries += 1
                self._pool_stats["connection_errors"] += 1
                self._record_connection_metric("error", time.time() - start_time, error=str(e))
                
                if retries > self._retry_policy["max_retries"]:
                    logger.error(f"获取数据库会话失败，已达最大重试次数: {str(e)}")
                    raise
                
                # 计算重试延迟（指数退避）
                delay = self._retry_policy["retry_delay"] * (self._retry_policy["backoff_factor"] ** (retries - 1))
                logger.warning(f"获取数据库会话失败，{delay:.2f}秒后重试 ({retries}/{self._retry_policy['max_retries']}): {str(e)}")
                await asyncio.sleep(delay)
    
    async def execute_with_connection(self, func: Callable[[AsyncConnection], Any]) -> Any:
        """
        直接使用数据库连接执行操作（适用于不需要ORM的场景）
        """
        if not self.engine:
            raise RuntimeError("数据库引擎未初始化")
        
        async with self.engine.connect() as conn:
            try:
                self._increment_connection_count()
                return await func(conn)
            finally:
                self._decrement_connection_count()
    
    async def execute_query(self, query: str, params: Optional[Dict[str, Any]] = None) -> Any:
        """
        直接执行SQL查询
        """
        async def _execute(conn):
            result = await conn.execute(text(query), params or {})
            return result
        
        return await self.execute_with_connection(_execute)
    
    def get_pool_stats(self) -> Dict[str, Any]:
        """
        获取连接池统计信息
        """
        return {
            **self._pool_stats,
            "last_stats_update": self._pool_stats["last_stats_update"].isoformat(),
            "pool_size": self._last_pool_size,
        }
    
    def _start_health_check(self):
        """
        启动连接池健康检查任务
        """
        if self._health_check_task:
            return
        
        async def _health_check_loop():
            while True:
                try:
                    await self._perform_health_check()
                    
                    # 根据负载情况动态调整连接池大小
                    if self._pool_adjustment_enabled:
                        await self._adjust_pool_size()
                    
                except Exception as e:
                    logger.error(f"连接池健康检查失败: {str(e)}")
                
                # 等待下一次检查
                await asyncio.sleep(self._health_check_interval)
        
        self._health_check_task = asyncio.create_task(_health_check_loop())
        logger.info("数据库连接池健康检查已启动")
    
    async def _perform_health_check(self):
        """
        执行连接池健康检查
        """
        if not self.engine:
            return
        
        try:
            start_time = time.time()
            
            # 执行简单的ping查询
            async with self.engine.connect() as conn:
                await conn.execute(text("SELECT 1"))
                
            response_time = time.time() - start_time
            
            # 更新连接池统计信息
            self._update_pool_metrics()
            
            logger.debug(f"数据库连接池健康检查成功，响应时间: {response_time:.4f}秒")
            
        except Exception as e:
            logger.error(f"数据库连接池健康检查失败: {str(e)}")
    
    def _update_pool_metrics(self):
        """
        更新连接池指标
        """
        # 注：SQLAlchemy不直接提供这些信息，这里使用估算值
        # 在生产环境中，可以考虑使用prometheus等监控系统获取更准确的指标
        self._pool_stats["last_stats_update"] = datetime.now()
        
        # 清理过期的连接指标（保留最近5分钟的数据）
        cutoff_time = datetime.now() - timedelta(minutes=5)
        for metric_type in self._connection_metrics:
            self._connection_metrics[metric_type] = [
                m for m in self._connection_metrics[metric_type]
                if m["timestamp"] > cutoff_time
            ]
    
    async def _adjust_pool_size(self):
        """
        基于使用情况动态调整连接池大小
        注意：SQLAlchemy的连接池大小在初始化后无法直接调整
        这里仅作为监控和建议，实际调整需要重启应用或重新初始化连接池
        """
        # 实现动态调整逻辑（示例）
        active_ratio = 0
        if self._pool_stats["total_connections"] > 0:
            active_ratio = self._pool_stats["active_connections"] / self._pool_stats["total_connections"]
        
        # 分析最近的连接错误率
        total_attempts = self._pool_stats["connection_successes"] + self._pool_stats["connection_errors"]
        error_rate = 0
        if total_attempts > 0:
            error_rate = self._pool_stats["connection_errors"] / total_attempts
        
        # 记录分析结果（实际调整需要重新初始化连接池）
        if error_rate > 0.1:
            logger.warning(f"数据库连接错误率较高 ({error_rate:.2%})，建议检查数据库负载和网络连接")
        
        if active_ratio > 0.8:
            logger.info(f"数据库连接池使用率较高 ({active_ratio:.2%})，可能需要考虑增加pool_size")
    
    def _increment_connection_count(self):
        """
        增加活动连接计数
        """
        self._pool_stats["active_connections"] += 1
        self._pool_stats["total_connections"] += 1
    
    def _decrement_connection_count(self):
        """
        减少活动连接计数
        """
        if self._pool_stats["active_connections"] > 0:
            self._pool_stats["active_connections"] -= 1
            self._pool_stats["idle_connections"] = self._pool_stats["total_connections"] - self._pool_stats["active_connections"]
    
    def _record_connection_metric(self, metric_type: str, duration: float, error: Optional[str] = None):
        """
        记录连接使用指标
        """
        if metric_type not in self._connection_metrics:
            self._connection_metrics[metric_type] = []
        
        self._connection_metrics[metric_type].append({
            "timestamp": datetime.now(),
            "duration": duration,
            "error": error,
        })
    
    async def test_connection(self) -> bool:
        """
        测试数据库连接是否正常
        """
        try:
            await self.execute_query("SELECT 1")
            return True
        except Exception as e:
            logger.error(f"数据库连接测试失败: {str(e)}")
            return False


# 创建全局异步数据库连接池管理器实例
async_db_pool_manager = AsyncDBPoolManager()


async def get_async_db_session() -> AsyncGenerator[AsyncSession, None]:
    """
    获取异步数据库会话（FastAPI依赖注入使用）
    """
    async for session in async_db_pool_manager.get_db_session():
        yield session


async def initialize_async_db_pool():
    """
    初始化异步数据库连接池
    """
    await async_db_pool_manager.initialize()


async def close_async_db_pool():
    """
    关闭异步数据库连接池
    """
    await async_db_pool_manager.close()
